from spec import plotlines
from astro.utilities import convolve_psf
import matplotlib.pyplot as pl
import numpy as np
from matplotlib import rc

A4LANDSCAPE = 11.7, 8.3
A4PORTRAIT = 8.3, 11.7

# SDSS galaxy line list
sdss_lines = """\
#1033.30  OVI
1215.67 1 Ly_alpha
#1239.42  NV
#1305.53  OI
#1335.52  CII
#1399.8   SiIV+OIV  
1545.86 1 CIV
#1640.4   HeII
#1665.85  OIII
#1857.4   AlIII
1908.27 1 CIII
#2326.0   CII
#2439.5   NeIV
2800.32  1 MgII
#3346.79  NeV
#3426.85  NeV
3728.30  1 [OII]
#3798.976 H-6
#3836.47  H-5
#3889.0   HeI
3934.777 2 K
3969.588 2 H
#4072.3   SII
4102.89  2 H$\\delta$
4305.61  2 G
4341.68  2 H$\\gamma$
4364.436 1 [OIII]
4862.68  1 H$\\beta$
4960.295 1 [OIII]
5008.240 1 [OIII]
5176.7   2 Mg
#5895.6   Na
#6302.046 1 OI
#6365.536 OI
6549.86  1 [NII]
6564.61  1 H$\\alpha$
6585.27  1 [NII]
#6707.89  Li
6718.29  1 [SII]
6732.67  1 [SII]
"""
temp = []
for l in sdss_lines.split('\n'):
    if not l.strip() or l.lstrip()[0] == '#':
        continue
    l = l.split()
    wa = float(l[0])
    ion = l[2]
    temp.append((ion, wa))

LINES = np.rec.fromrecords(temp, names='ion,wa')


class PlotLines(object):
    """ Add a keypress function callback to a figure, that plots lines
    at a given redshift.

    Remove the callback using fig.canvas.mpl_disconnect(self.cid).

    TODO: confine to a particular axis in a figure?
    """
    
    def __init__(self, fig, z=None, lines=None, disconnect_other=True):
        """ Lines is a structured array of transitions with fields wa
        and ion.
        """
        keys = '1 2 3 4 f1 f2 f3 f4 5 6 7 8 9 0 f5 f6 f7 f8'.split()
        if lines is None:
            lines = LINES
        lines1 = [(k,i,w) for k,i,w in zip(keys, lines['ion'], lines['wa'])]
        temp = []
        for key,ion,wa in lines1:
            wa = str(wa).split('.')[0]
            temp.append('%s: %s %s\n' % (key, ion, wa))
        temp.append('\n?: Print this message\n')
        temp.append('q: Quit\n')
        
        self.help = ''.join(temp)
        #self.ax = ax
        self.fig = fig
        self.lines = lines
        self.keys = [l[0] for l in lines1]
        if disconnect_other:
            # disable any existing key press callbacks
            cids = list(fig.canvas.callbacks.callbacks['key_press_event'])
            for cid in cids:
                fig.canvas.callbacks.disconnect(cid)
        
        self.cid = self.fig.canvas.mpl_connect('key_press_event',
                                               self.on_keypress)
        self.artists = None
        if z is not None:
            self.plotlines(z+1, pl.gca())
        print self.help
        self.wait = True
        self.z = None
        self.zgood = None

    def plotlines(self, zp1, ax):
        if self.artists:
            for artist in self.artists:
                try:
                    artist.remove()
                except ValueError:
                    # axes have been removed since we added these lines
                    break
        self.artists = plotlines(zp1-1, ax, lines=self.lines)

    def on_keypress(self, event):
        if event.key == 'q':
            self.disconnect()
            return
        elif event.key == '?':
            print self.help
            return
        elif event.key == 'k':
            print 'keeping redshift'
            self.zgood = self.z
            return
        elif event.inaxes is None:
            return
        try:
            i = self.keys.index(event.key)
        except ValueError:
            return
        line = self.lines[i]
        zp1 = event.xdata / line['wa']
        self.z = zp1 - 1
        print 'z=%.3f, %s %f' % (zp1 - 1, line['ion'], line['wa'])
        self.plotlines(zp1, event.inaxes)
        self.fig.canvas.draw()
        
    def disconnect(self):
        self.fig.canvas.mpl_disconnect(self.cid)
        self.wait = False


class Smooth(object):
    """ Adds functionality to plot smoothed flux.

    keys used: s
    """
    def __init__(self, fig, artist, disconnect_other=True):
      
        #self.ax = ax
        self.fig = fig
        self.fl = artist.get_ydata()
        self.sfl = self.fl
        self.artist = artist
        if disconnect_other:
            # disable any existing key press callbacks
            cids = list(fig.canvas.callbacks.callbacks['key_press_event'])
            for cid in cids:
                fig.canvas.callbacks.disconnect(cid)
        
        self.cid = self.fig.canvas.mpl_connect('key_press_event',
                                               self.on_keypress)
        self.help = """\
s: smooth spectrum by convolution with a Gaussian.
"""
        print self.help
        self.wait = True
      
    def on_keypress(self, event):
        if event.key != 's':
            return
        c = raw_input('FWHM in pixels of Gaussian to convolve with? '
                      '(blank to restore spectrum) ')
        if c == '':
            # restore spectrum
            self.sfl = self.fl
        else:
            try:
                fwhm = int(c)
            except TypeError:
                print 'FWHM must be an integer'
            if fwhm == 0:
                self.sfl = self.fl
            else:
                self.sfl = convolve_psf(self.fl, fwhm, edge='reflect')

        self.artist.set_ydata(self.sfl)
        self.fig.canvas.draw()
                
    def disconnect(self):
        self.fig.canvas.mpl_disconnect(self.cid)
        self.wait = False

class GridPlots(object):
    def __init__(self, rows=10, cols=1, figsize=A4LANDSCAPE):
        rc('xtick.minor', size=2)
        rc('ytick.minor', size=2)
        rc('xtick.major', size=4)
        rc('ytick.major', size=4)
        fig = pl.figure(figsize=figsize)
        x0, xn = 0.1, 0.95
        y0, yn = 0.1, 0.9
        height = (yn - y0) / float(rows)
        width = (xn - x0) / float(cols)
        axes = []
        for i in range(rows):
            axcols = []
            for j in range(cols):
                axcols.append(fig.add_axes((x0 + j * width, y0 + i * height,
                                            width, height)))
            axes.append(axcols)
        # go from top to bottom
        self.ax = np.array(axes)
        self.ax = self.ax[::-1]
        #self.ax.squeeze()
        for row in self.ax[:-1]:
            for ax in row:
                ax.set_xticklabels([])
        for col in self.ax[:, 1:]:
            for ax in col:
                ax.set_yticklabels([])
        rc('defaults')
        
    def cleanup(self):
        for row in self.ax[:-1]:
            for ax in row:
                ax.set_xticklabels([])
        for col in self.ax[:, 1:]:
            for ax in col:
                ax.set_yticklabels([])
        pl.show()


def example():
    # fit a redshift by eye interactively
    f = pl.gcf()
    # add a spectrum to the figure, then wrap with:
    temp = PlotLines(f)
    # loop where we play around with different redshifts interactively
    while temp.wait:
        pl.waitforbuttonpress()

    redshift = temp.zgood
        


